// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include <pollux/dwio/common/encryption/test_provider.h>
#include <tests/dwio/dwrf/orc_test.h>
#include <pollux/dwio/dwrf/utils/proto_utils.h>
#include <pollux/type/hive/hive_type_parser.h>

using namespace kumo::pollux::dwio::common;
using namespace kumo::pollux::dwio::common::encryption;
using namespace kumo::pollux::dwio::common::encryption::test;
using namespace kumo::pollux::dwrf;
using namespace kumo::pollux::dwrf::encryption;
using namespace kumo::pollux::type::hive;

TEST(Decryption, NotEncrypted) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  TestDecrypterFactory factory;
  auto handler = DecryptionHandler::create(footer, &factory);
  ASSERT_FALSE(handler->isEncrypted());
}

TEST(Decryption, NoKeyProvider) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  footer.mutable_encryption();
  TestDecrypterFactory factory;
  ASSERT_THROW(
      DecryptionHandler::create(footer, &factory), exception::LoggedException);
}

TEST(Decryption, EmptyGroup) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);
  TestDecrypterFactory factory;
  ASSERT_THROW(
      DecryptionHandler::create(footer, &factory), exception::LoggedException);
}

TEST(Decryption, EmptyNodes) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);
  auto group = enc->add_encryptiongroups();
  group->set_keymetadata("key");
  TestDecrypterFactory factory;
  ASSERT_THROW(
      DecryptionHandler::create(footer, &factory), exception::LoggedException);
}

TEST(Decryption, StatsMismatch) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);
  auto group = enc->add_encryptiongroups();
  group->set_keymetadata("key");
  group->add_nodes(1);
  group->add_nodes(2);
  group->add_statistics();
  TestDecrypterFactory factory;
  ASSERT_THROW(
      DecryptionHandler::create(footer, &factory), exception::LoggedException);
}

TEST(Decryption, KeyExistenceMismatch) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);
  for (size_t i = 0; i < 2; ++i) {
    auto group = enc->add_encryptiongroups();
    if (i == 0) {
      group->set_keymetadata("key");
    }
    group->add_nodes(i + 1);
    group->add_statistics();
  }
  TestDecrypterFactory factory;
  ASSERT_THROW(
      DecryptionHandler::create(footer, &factory), exception::LoggedException);
}

TEST(Decryption, ReuseStripeKey) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);
  auto group = enc->add_encryptiongroups();
  group->add_nodes(1);
  group->add_statistics();
  auto stripe = footer.add_stripes();
  *stripe->add_keymetadata() = "foobar";
  TestDecrypterFactory factory;
  auto handler = DecryptionHandler::create(footer, &factory);
  auto& td =
      dynamic_cast<const TestEncryption&>(handler->getEncryptionProvider(1));
  ASSERT_EQ(td.getKey(), "foobar");
}

TEST(Decryption, StripeKeyMismatch) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);
  auto group = enc->add_encryptiongroups();
  group->add_nodes(1);
  group->add_statistics();
  auto stripe = footer.add_stripes();
  *stripe->add_keymetadata() = "foobar";
  *stripe->add_keymetadata() = "foobar";
  TestDecrypterFactory factory;
  ASSERT_THROW(
      DecryptionHandler::create(footer, &factory), exception::LoggedException);
}

TEST(Decryption, Basic) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int,b:float,c:string,d:double>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);
  for (auto i = 0; i < 5; ++i) {
    if (i % 2 == 1) {
      auto group = enc->add_encryptiongroups();
      group->add_nodes(i);
      group->set_keymetadata(melon::to<std::string>("key", i));
      group->add_statistics();
    }
  }
  TestDecrypterFactory factory;
  auto handler = DecryptionHandler::create(footer, &factory);
  ASSERT_TRUE(handler->isEncrypted());
  for (auto i = 0; i < 5; ++i) {
    ASSERT_EQ(handler->isEncrypted(i), (i % 2 == 1));
    if (i % 2 == 1) {
      ASSERT_EQ(
          dynamic_cast<const TestEncryption&>(handler->getEncryptionProvider(i))
              .getKey(),
          melon::to<std::string>("key", i));
    }
  }
}

TEST(Decryption, NestedType) {
  HiveTypeParser parser;
  auto type = parser.parse(
      "struct<a:int,b:map<float,map<int,double>>,c:struct<a:string,b:int>,d:array<double>>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);

  uint32_t nodeCount = 12;
  std::vector<uint32_t> cols{1, 2, 7, 10, nodeCount};
  std::vector<uint32_t> encRoots{1, 2, 9, 10, nodeCount};
  melon::F14FastMap<uint32_t, uint32_t> nodeToRoot;
  for (auto i = 0; i < cols.size() - 1; ++i) {
    auto from = std::max(cols[i], encRoots[i]);
    auto to = std::min(cols[i + 1], encRoots[i + 1]);
    for (auto j = from; j < to; ++j) {
      nodeToRoot[j] = from;
    }
  }

  for (auto i = 0; i < encRoots.size() - 1; ++i) {
    auto group = enc->add_encryptiongroups();
    group->add_nodes(encRoots[i]);
    group->set_keymetadata(melon::to<std::string>("key", i));
    group->add_statistics();
  }
  TestDecrypterFactory factory;
  auto handler = DecryptionHandler::create(footer, &factory);
  ASSERT_TRUE(handler->isEncrypted());
  ASSERT_EQ(handler->getKeyProviderType(), EncryptionProvider::Unknown);
  ASSERT_EQ(handler->getEncryptionGroupCount(), 4);
  for (auto i = 0; i < nodeCount; ++i) {
    ASSERT_EQ(handler->isEncrypted(i), nodeToRoot.count(i) > 0);
    if (handler->isEncrypted(i)) {
      ASSERT_EQ(handler->getEncryptionRoot(i), nodeToRoot[i]);
    }
  }
}

TEST(Decryption, RootNode) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:int,b:int>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);
  auto group = enc->add_encryptiongroups();
  group->add_nodes(0);
  group->set_keymetadata("key");
  group->add_statistics();
  TestDecrypterFactory factory;
  auto handler = DecryptionHandler::create(footer, &factory);
  ASSERT_EQ(handler->getEncryptionGroupCount(), 1);
}

TEST(Decryption, GroupOverlap) {
  HiveTypeParser parser;
  auto type = parser.parse("struct<a:struct<a:float,b:double>>");
  proto::Footer footer;
  ProtoUtils::writeType(*type, footer);
  auto enc = footer.mutable_encryption();
  enc->set_keyprovider(proto::Encryption_KeyProvider_UNKNOWN);

  auto group1 = enc->add_encryptiongroups();
  group1->add_nodes(1);
  group1->set_keymetadata("key");
  group1->add_statistics();

  auto group2 = enc->add_encryptiongroups();
  group2->add_nodes(3);
  group2->set_keymetadata("key");
  group2->add_statistics();

  TestDecrypterFactory factory;
  ASSERT_THROW(
      DecryptionHandler::create(footer, &factory), exception::LoggedException);
}
